import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
from functools import partial
import matplotlib.pyplot as plt
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import threading
import time
import logging
import gc

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 常量
BATCH_SIZE = 200  # 增加批次大小
CONSECUTIVE_FAILURE_THRESHOLD = 3  # 连续失败阈值
MAX_WORKERS = max(2, multiprocessing.cpu_count() - 1)  # 保留一个CPU核心给系统
RETRY_THRESHOLD = 0.05  # 单批次失败率阈值降低到5%

class EnhancedImageStabilizer:
    def __init__(self):
        self.reference_frame = None
        self.reference_right_edge = None
        self.crop_width = None
        self.roi = None  # 存储用户标注的感兴趣区域
        self.consecutive_failures = 0  # 连续检测失败计数

    def set_roi(self, x: int, y: int, w: int, h: int):
        """设置用户标注的感兴趣区域"""
        self.roi = (x, y, w, h)
        logging.info(f"已设置ROI区域: x={x}, y={y}, width={w}, height={h}")
        # 重置连续失败计数
        self.consecutive_failures = 0

    def detect_vertical_lines(self, gray_img: np.ndarray) -> np.ndarray:
        """检测图像中的垂直线条，改进版本"""
        # 如果设置了ROI，只处理ROI区域
        if self.roi is not None:
            x, y, w, h = self.roi
            roi_img = gray_img[y:y+h, x:x+w]
        else:
            roi_img = gray_img

        # 优化对比度增强参数
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(16,16))
        enhanced = clahe.apply(roi_img)

        # 使用更快的Sobel边缘检测替代Canny
        sobel_x = cv2.Sobel(enhanced, cv2.CV_64F, 1, 0, ksize=3)
        sobel_x = np.absolute(sobel_x)
        sobel_x = np.uint8(sobel_x)
        
        # 二值化
        _, binary = cv2.threshold(sobel_x, 30, 255, cv2.THRESH_BINARY)

        # 使用更小的kernel进行形态学操作
        kernel_vertical = np.ones((11,1), np.uint8)
        vertical_edges = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)

        return vertical_edges

    def find_right_edge(self, binary_img: np.ndarray) -> int | None:
        """在二值化图像中找到右边缘，改进版本"""
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)

        # 使用一维高斯卷积减少噪声影响
        window = 15
        sigma = window / 3.0
        x = np.linspace(-window // 2, window // 2, window)
        gaussian = np.exp(-(x**2) / (2 * sigma**2))
        gaussian = gaussian / np.sum(gaussian)
        smoothed = np.convolve(projection, gaussian, mode='valid')

        # 计算梯度
        gradient = np.gradient(smoothed)

        # 使用更强的阈值检测边缘
        threshold = np.std(gradient) * 2.5

        # 从右向左扫描找到第一个显著边缘
        right_edge = None
        for i in range(len(gradient) - 1, window, -1):
            if abs(gradient[i]) > threshold:
                # 使用局部最大值优化边缘位置
                local_region = gradient[max(0, i - 8):min(len(gradient), i + 9)]
                local_offset = np.argmax(np.abs(local_region))
                right_edge = i - 8 + local_offset
                break

        # 如果没有找到显著边缘，尝试使用更宽的搜索窗口和更低的阈值
        if right_edge is None:
            threshold = np.std(gradient) * 1.5
            for i in range(len(gradient) - 1, window, -1):
                if abs(gradient[i]) > threshold:
                    local_region = gradient[max(0, i - 10):min(len(gradient), i + 11)]
                    local_offset = np.argmax(np.abs(local_region))
                    right_edge = i - 10 + local_offset
                    break

        # 如果设置了ROI，需要将检测到的边缘位置转换回原图坐标系
        if self.roi is not None and right_edge is not None:
            right_edge += self.roi[0]

        return right_edge

    def initialize_with_reference(self, reference_img: np.ndarray) -> bool:
        """使用参考帧初始化稳定器"""
        self.reference_frame = reference_img.copy()

        # 显示参考帧，让用户标注感兴趣区域
        cv2.imshow("标注参考帧", reference_img)
        logging.info("请在图像上标注棒束通道和胶布的交界区域（矩形框）")
        logging.info("按下鼠标左键开始绘制，松开鼠标左键完成绘制")
        logging.info("按's'键保存标注，按'q'键退出")

        # 等待用户标注
        self.roi = cv2.selectROI("标注参考帧", reference_img, False)
        cv2.destroyAllWindows()

        if self.roi == (0, 0, 0, 0):
            logging.warning("警告：未标注区域，将使用全图检测")
            self.roi = None
        else:
            logging.info(f"已标注ROI区域: {self.roi}")

        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)

        # 检测垂直线
        binary = self.detect_vertical_lines(gray)

        # 找到右边缘
        right_edge = self.find_right_edge(binary)

        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的固定距离）
            self.crop_width = right_edge - 100  # 可以根据需要调整这个值
            logging.info(f"初始化完成：参考右边缘位置: {right_edge}")
            logging.info(f"裁剪宽度设置为: {self.crop_width}")
            return True
        else:
            logging.warning("警告：无法在参考帧中检测到右边缘！")
            return False

    def stabilize_image(self, img: np.ndarray) -> tuple[np.ndarray, bool]:
        """稳定图像，通过检测右边缘并调整裁剪区域"""
        if self.reference_right_edge is None or self.crop_width is None:
            return img, False

        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)

        # 检测垂直线
        binary = self.detect_vertical_lines(gray)

        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)

        if current_right_edge is None:
            logging.warning("警告：当前帧未检测到右边缘，使用参考帧的裁剪参数")
            # 使用参考帧的裁剪参数
            start_x = max(0, self.reference_right_edge - self.crop_width)
            end_x = self.reference_right_edge
            cropped = img[:, start_x:end_x]
            return cropped, False
        else:
            # 计算裁剪区域，使用动态边缘位置
            start_x = max(0, current_right_edge - self.crop_width)
            end_x = current_right_edge
            cropped = img[:, start_x:end_x]
            return cropped, True

def process_single_image(args: tuple[str, EnhancedImageStabilizer]) -> tuple[np.ndarray | None, bool, str]:
    """处理单张图片"""
    img_path, cropper = args
    try:
        # 读取当前帧
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)

        if current_frame is None:
            return None, False, os.path.basename(img_path)

        # 裁剪图像并获取检测状态
        cropped_frame, detection_success = cropper.stabilize_image(current_frame)

        # 返回处理后的帧、检测状态和文件名
        return cropped_frame, detection_success, os.path.basename(img_path)

    except Exception as e:
        logging.error(f"处理图片 {img_path} 时发生错误: {str(e)}")
        return None, False, os.path.basename(img_path)

def process_images_to_images(input_path: str, output_path: str, progress_callback=None, status_callback=None, cancel_flag: threading.Event = None, output_format: str = "bmp"):
    """处理图片序列"""
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        total_images = len(image_files)
        
        if total_images == 0:
            if status_callback:
                status_callback("未找到BMP文件！")
            return

        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)

        # 读取第一帧作为参考帧
        first_frame_path = os.path.join(input_path, image_files[0])
        reference_frame = cv2.imdecode(np.fromfile(first_frame_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            if status_callback:
                status_callback(f"无法读取参考帧：{image_files[0]}")
            return

        # 初始化图像稳定器
        cropper = EnhancedImageStabilizer()
        if not cropper.initialize_with_reference(reference_frame):
            if status_callback:
                status_callback("初始化失败！")
            return

        # 创建进程池
        with Pool(processes=MAX_WORKERS) as pool:
            processed_count = 0
            failed_images = []
            consecutive_failures = 0
            i = 0
            
            while i < len(image_files):
                if cancel_flag and cancel_flag.is_set():
                    if status_callback:
                        status_callback("处理已取消")
                    break
                
                # 准备当前批次
                batch_files = image_files[i:i + BATCH_SIZE]
                image_paths = [os.path.join(input_path, img) for img in batch_files]
                process_args = [(path, cropper) for path in image_paths]
                
                # 处理当前批次
                results = pool.map(process_single_image, process_args)
                
                # 统计当前批次的失败数
                batch_failures = 0
                batch_results = []
                
                for result in results:
                    cropped_frame, detection_success, filename = result
                    if cropped_frame is not None:
                        # 保存处理后的图片
                        output_filename = os.path.splitext(filename)[0] + f".{output_format}"
                        output_file = os.path.join(output_path, output_filename)
                        cv2.imencode(f".{output_format}", cropped_frame)[1].tofile(output_file)
                        
                        if not detection_success:
                            batch_failures += 1
                            failed_images.append(filename)
                    else:
                        batch_failures += 1
                        failed_images.append(filename)
                    
                    batch_results.append((cropped_frame, detection_success, filename))
                    processed_count += 1
                
                # 计算当前批次的失败率
                failure_rate = batch_failures / len(batch_results)
                
                # 如果当前批次失败率高于阈值或连续失败次数达到阈值
                if failure_rate > RETRY_THRESHOLD or consecutive_failures >= CONSECUTIVE_FAILURE_THRESHOLD:
                    if status_callback:
                        status_callback(f"\n当前批次失败率: {failure_rate*100:.1f}% 或连续失败次数: {consecutive_failures}")
                        status_callback("需要重新标注ROI区域...")
                    
                    # 读取当前批次的第一帧用于重新标注
                    current_frame = cv2.imdecode(
                        np.fromfile(image_paths[0], dtype=np.uint8),
                        cv2.IMREAD_COLOR
                    )
                    
                    # 显示图片并等待用户标注
                    cv2.imshow("重新标注ROI区域", current_frame)
                    new_roi = cv2.selectROI("重新标注ROI区域", current_frame, False)
                    cv2.destroyAllWindows()
                    
                    if new_roi != (0, 0, 0, 0):
                        cropper.set_roi(*new_roi)
                        if status_callback:
                            status_callback("已更新ROI区域，重新处理当前批次...")
                        consecutive_failures = 0
                        continue  # 重新处理当前批次
                
                # 更新连续失败计数
                if failure_rate > RETRY_THRESHOLD:
                    consecutive_failures += 1
                else:
                    consecutive_failures = 0
                
                # 更新进度
                if progress_callback:
                    progress_callback(processed_count / total_images * 100)
                
                # 移动到下一批次
                i += BATCH_SIZE
                
                # 强制垃圾回收
                gc.collect()
            
            # 处理完成后的报告
            if status_callback:
                status_message = f"处理完成！共处理 {processed_count} 张图片。"
                if failed_images:
                    status_message += f"\n{len(failed_images)} 张图片检测失败："
                    status_message += "\n" + ", ".join(failed_images[:10])
                    if len(failed_images) > 10:
                        status_message += f"... 等 {len(failed_images)} 张"
                status_callback(status_message)
                
    except Exception as e:
        if status_callback:
            status_callback(f"处理过程中发生错误：{str(e)}")
        logging.error(f"处理过程中发生错误：{str(e)}")

class PIVProcessorApp:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("PIV图像处理工具")
        self.root.geometry("600x400")

        # 创建主框架
        self.main_frame = ttk.Frame(root, padding="20")
        self.main_frame.pack(fill=tk.BOTH, expand=True)

        # 输入文件夹选择
        input_frame = ttk.Frame(self.main_frame)
        input_frame.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(input_frame, text="输入文件夹:", font=("Arial", 12)).pack(side=tk.LEFT, padx=(0, 10))
        self.input_path = tk.StringVar()
        ttk.Entry(input_frame, textvariable=self.input_path, width=40, font=("Arial", 10)).pack(side=tk.LEFT, expand=True, fill=tk.X)
        ttk.Button(input_frame, text="浏览...", command=self.select_input_folder, style="TButton").pack(side=tk.LEFT, padx=(10, 0))

        # 输出文件夹选择
        output_frame = ttk.Frame(self.main_frame)
        output_frame.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(output_frame, text="输出文件夹:", font=("Arial", 12)).pack(side=tk.LEFT, padx=(0, 10))
        self.output_path = tk.StringVar()
        ttk.Entry(output_frame, textvariable=self.output_path, width=40, font=("Arial", 10)).pack(side=tk.LEFT, expand=True, fill=tk.X)
        ttk.Button(output_frame, text="浏览...", command=self.select_output_folder, style="TButton").pack(side=tk.LEFT, padx=(10, 0))

        # 输出格式选择
        format_frame = ttk.Frame(self.main_frame)
        format_frame.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(format_frame, text="输出格式:", font=("Arial", 12)).pack(side=tk.LEFT, padx=(0, 10))
        self.output_format = tk.StringVar(value="bmp")  # 默认格式为bmp
        output_formats = ["bmp", "png", "jpg"]
        self.format_combobox = ttk.Combobox(format_frame, textvariable=self.output_format, values=output_formats, width=5, font=("Arial", 10), state="readonly")
        self.format_combobox.pack(side=tk.LEFT, padx=(0, 10))

        # 进度条
        progress_frame = ttk.Frame(self.main_frame)
        progress_frame.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(progress_frame, text="处理进度:", font=("Arial", 12)).pack(side=tk.LEFT, padx=(0, 10))
        self.progress = ttk.Progressbar(progress_frame, orient=tk.HORIZONTAL, length=100, mode='determinate')
        self.progress.pack(side=tk.LEFT, expand=True, fill=tk.X)

        # 状态信息
        status_frame = ttk.Frame(self.main_frame)
        status_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        ttk.Label(status_frame, text="状态:", font=("Arial", 12)).pack(side=tk.LEFT, padx=(0, 10))
        self.status_text = tk.Text(status_frame, height=10, width=50, wrap=tk.WORD, font=("Arial", 10))
        self.status_text.pack(side=tk.LEFT, expand=True, fill=tk.BOTH)
        scrollbar = ttk.Scrollbar(status_frame, orient=tk.VERTICAL, command=self.status_text.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.status_text.configure(yscrollcommand=scrollbar.set)

        # 按钮
        button_frame = ttk.Frame(self.main_frame)
        button_frame.pack(fill=tk.X, pady=(10, 0))
        self.start_button = ttk.Button(button_frame, text="开始处理", command=self.start_processing, style="TButton")
        self.start_button.pack(side=tk.RIGHT, padx=(0, 10))
        self.cancel_button = ttk.Button(button_frame, text="取消处理", command=self.cancel_processing, style="TButton")
        self.cancel_button.pack(side=tk.RIGHT, padx=(0, 10))
        self.reset_button = ttk.Button(button_frame, text="重置", command=self.reset_settings, style="TButton")  # 添加重置按钮
        self.reset_button.pack(side=tk.RIGHT, padx=(0, 10))

        # 设置样式
        style = ttk.Style()
        style.configure("TButton", font=("Arial", 10))
        style.configure("TLabel", font=("Arial", 12))

        # 处理线程
        self.processing_thread = None
        self.cancel_flag = threading.Event()  # 取消标志

    def reset_settings(self):
        """重置设置"""
        self.input_path.set("")
        self.output_path.set("")
        self.output_format.set("bmp")
        self.progress['value'] = 0
        self.status_text.delete(1.0, tk.END)
        self.start_button.configure(state="normal")
        self.cancel_button.configure(state="normal")
        if self.processing_thread and self.processing_thread.is_alive():
            self.cancel_processing() # 如果线程还在运行，先取消

    def select_input_folder(self):
        """选择输入文件夹"""
        folder = filedialog.askdirectory(title="选择输入文件夹")
        if folder:
            self.input_path.set(folder)

    def select_output_folder(self):
        """选择输出文件夹"""
        folder = filedialog.askdirectory(title="选择输出文件夹")
        if folder:
            self.output_path.set(folder)

    def update_progress(self, value: float):
        """更新进度条"""
        self.progress['value'] = value
        self.root.update_idletasks()

    def update_status(self, message: str):
        """更新状态信息"""
        self.status_text.insert(tk.END, message + "\n")
        self.status_text.see(tk.END)
        self.root.update_idletasks()

    def start_processing(self):
        """开始处理"""
        input_path = self.input_path.get()
        output_path = self.output_path.get()
        output_format = self.output_format.get()

        if not input_path or not output_path:
            messagebox.showerror("错误", "请选择输入和输出文件夹")
            return

        if not os.path.exists(input_path):
            messagebox.showerror("错误", f"输入文件夹 {input_path} 不存在")
            return

        # 禁用开始按钮和取消按钮
        self.start_button.configure(state="disabled")
        self.cancel_button.configure(state="disabled")

        # 重置进度条
        self.progress['value'] = 0

        # 清空状态文本
        self.status_text.delete(1.0, tk.END)
        self.update_status(f"开始处理图片...")
        self.update_status(f"输入文件夹: {input_path}")
        self.update_status(f"输出文件夹: {output_path}")
        self.update_status(f"输出格式: {output_format}")

        # 创建并启动处理线程
        self.processing_thread = threading.Thread(
            target=self.process_images_thread,
            args=(input_path, output_path)
        )
        self.processing_thread.daemon = True
        self.processing_thread.start()

    def process_images_thread(self, input_path: str, output_path: str):
        """处理图像线程"""
        try:
            self.cancel_flag.clear()  # 清除取消标志
            success = process_images_to_images(
                input_path,
                output_path,
                progress_callback=lambda value: self.root.after(0, self.update_progress, value),
                status_callback=lambda msg: self.root.after(0, self.update_status, msg),
                cancel_flag=self.cancel_flag,
                output_format=self.output_format.get()
            )

            if success:
                self.root.after(0, lambda: messagebox.showinfo("完成", "图像处理完成！"))
            else:
                self.root.after(0, lambda: messagebox.showerror("错误", "处理过程中发生错误，请查看状态信息"))
        except Exception as e:
            self.root.after(0, lambda: messagebox.showerror("错误", f"处理过程中发生异常: {str(e)}"))
        finally:
            # 重新启用开始按钮
            self.root.after(0, lambda: self.start_button.configure(state="normal"))
            # 重新启用取消按钮
            self.root.after(0, lambda: self.cancel_button.configure(state="normal"))


    def cancel_processing(self):
        """取消处理"""
        if self.processing_thread and self.processing_thread.is_alive():
            self.cancel_flag.set()
            logging.info("正在尝试取消处理...")
            self.update_status("正在取消...")

if __name__ == '__main__':
    multiprocessing.freeze_support()  # 添加freeze_support()
    root = tk.Tk()
    app = PIVProcessorApp(root)
    root.mainloop()
